
#include "../include/qp_solver.hpp"


QPSolver::QPSolver(bool isShow) {
    this->show = isShow;
    solver.settings()->setWarmStart(true);
}

QPSolver::~QPSolver() {}


// 更新需要平滑的参考线
void QPSolver::updateRefePoint(vector<ReferencePoint> &ref_point) {
    this->ref_x.clear();
    this->ref_y.clear();
    for (int i = 0; i < ref_point.size(); i++) {
        this->ref_x.push_back(ref_point[i].x);
        this->ref_y.push_back(ref_point[i].y);
    }
}

void QPSolver::updateRefePoint(vector<TraPoint> &list) {
    this->ref_x.clear();
    this->ref_y.clear();
    for (int i = 0; i < list.size(); i++) {
        this->ref_x.push_back(list[i].x);
        this->ref_y.push_back(list[i].y);
    }
}

void QPSolver::updateRefePoint(vector<double> &x, vector<double> &y) {
    this->ref_x.clear();
    this->ref_y.clear();
    this->ref_x.assign(x.begin(), x.end());
    this->ref_y.assign(y.begin(), y.end());
}

/**
 * 初始化五次多项式求解矩阵
 * https://yunchengjiang.blog.csdn.net/article/details/113745314
 * https://blog.csdn.net/gophae/article/details/115122792
*/
void QPSolver::initQuinticPoly() {
    int NumberOfVariables = 6;   //A矩阵的列数
    int NumberOfConstraints = 6; //A矩阵的行数
    solver.data()->setNumberOfVariables(NumberOfVariables);
    solver.data()->setNumberOfConstraints(NumberOfConstraints);

    hessian.resize(6, 6);
    hessian1.resize(6, 6);
    hessian2.resize(6, 6);
    double Te = 5;
    double d0 = 1;
    double v0 = 0.2;
    double a0 = 0;
    double LB = 2.5;
    double UB = 3.5;

    for (int i = 0; i < 6; i++)
    {
        hessian1.insert(0, i) = pow(Te, (i + 1)) / (i + 1);
        hessian1.insert(1, i) = pow(Te, (i + 2)) / (i + 2);
        hessian1.insert(2, i) = pow(Te, (i + 3)) / (i + 3);
        hessian1.insert(3, i) = pow(Te, (i + 4)) / (i + 4);
        hessian1.insert(4, i) = pow(Te, (i + 5)) / (i + 5);
        hessian1.insert(5, i) = pow(Te, (i + 6)) / (i + 6);
    }

    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
        hessian2.insert(j, i) = 0;
        }
    }

    for (int i = 3; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
        hessian2.insert(i, j) = 0;
        }
    }
    hessian2.insert(3, 3) = 36 * Te;
    hessian2.insert(3, 4) = 72 * Te * Te;
    hessian2.insert(3, 5) = 120 * Te * Te * Te;
    hessian2.insert(4, 3) = 72 * Te * Te;
    hessian2.insert(4, 4) = 192 * Te * Te * Te;
    hessian2.insert(4, 5) = 360 * Te * Te * Te * Te;
    hessian2.insert(5, 3) = 120 * Te * Te * Te;
    hessian2.insert(5, 4) = 360 * Te * Te * Te * Te;
    hessian2.insert(5, 5) = 720 * Te * Te * Te * Te * Te;

    hessian = hessian1 + hessian2;
    // hessian.insert(0, 0) = 1;
    // hessian.insert(1, 0) = -1;
    // hessian.insert(0, 1) = -1;
    // hessian.insert(1, 1) = 2;
    std::cout
        << "hessian1:" << std::endl
        << hessian1 << std::endl
        << "hessian2: " << std::endl
        << hessian2 << std::endl;

    gradient.resize(6);
    gradient << 0, 0, 0, 0, 0, 0;

    std::cout << "gradient:" << std::endl;
    std::cout << gradient << std::endl;

    linearMatrix.resize(6, 6);
    // the first row
    for (int i = 0; i < 6; i++)
    {
        linearMatrix.insert(0, i) = pow(Te, i);
    }
    //the second row
    for (int i = 1; i < 6; i++)
    {
        linearMatrix.insert(1, i) = 0;
    }
    linearMatrix.insert(1, 0) = 1;

    // the third row
    for (int i = 2; i < 6; i++)
    {
        linearMatrix.insert(2, i) = 0;
    }
    linearMatrix.insert(2, 0) = 0;
    linearMatrix.insert(2, 1) = 1;
    // the fourth row
    linearMatrix.insert(3, 0) = 0;
    linearMatrix.insert(3, 1) = 1;
    linearMatrix.insert(3, 2) = 2 * Te;
    linearMatrix.insert(3, 3) = 3 * Te * Te;
    linearMatrix.insert(3, 4) = 4 * Te * Te * Te;
    linearMatrix.insert(3, 5) = 5 * Te * Te * Te * Te;
    // the fifth row
    linearMatrix.insert(4, 0) = 0;
    linearMatrix.insert(4, 1) = 0;
    linearMatrix.insert(4, 2) = 2;
    linearMatrix.insert(4, 3) = 0;
    linearMatrix.insert(4, 4) = 0;
    linearMatrix.insert(4, 5) = 0;
    // the sixth row
    linearMatrix.insert(5, 0) = 0;
    linearMatrix.insert(5, 1) = 0;
    linearMatrix.insert(5, 2) = 2;
    linearMatrix.insert(5, 3) = 6 * Te;
    linearMatrix.insert(5, 4) = 12 * Te * Te;
    linearMatrix.insert(5, 5) = 20 * Te * Te * Te;

    std::cout << "linearMatrix:" << std::endl
                << linearMatrix << std::endl;

    lowerBound.resize(6);
    // lowerBound << -OsqpEigen::INFTY, -OsqpEigen::INFTY, -OsqpEigen::INFTY;
    lowerBound << (d0 + LB), d0, v0, 0, a0, 0;
    std::cout << "lowerBound:" << std::endl;
    std::cout << lowerBound << std::endl;

    upperBound.resize(6);
    upperBound << (d0 + UB), d0, v0, 0, a0, 0;
    std::cout << "upperBound:" << std::endl;
    std::cout << upperBound << std::endl;
}


void QPSolver::initFEMSmoothMat(int n, double w_cost_smooth, double w_cost_length, double w_cost_ref, 
                                double lb, double ub) {
    solver.data()->setNumberOfVariables(2 * n);
    solver.data()->setNumberOfConstraints(2 * n);
    hessian.resize(2 * n, 2 * n);
    gradient.resize(2 * n);
    linearMatrix.resize(2 * n, 2 * n);
    lowerBound.resize(2 * n);
    upperBound.resize(2 * n);
    // A_1 = [1 0 -2 0 1 0
    //        0 1 0 -2 0 1
    //            1 0 -2 0 1 0
    //            0 1 0 -2 0 1
    //                ...........
    //
    // A_2 = [1 0 -1 0
    //        0 1 0 -1
    //            1 0 -1 0
    //            0 1 0 -1
    //                ........
    // A_3 为单位矩阵
    // 目前 reference line 的长度为 40 个点，所以求解矩阵规模为 80

    /****************** 稀疏矩阵不支持块操作 **********************/
    // MatrixXd A_1(2 * n - 4, 2 * n);
    // MatrixXd A_2(2 * n - 2, 2 * n);
    // MatrixXd A_3 = MatrixXd::Identity(2 * n, 2 * n);

    // matrix.block<p,q>(i,j); 以 i j 为起点，选取大小为 p q 的矩阵
    // MatrixXd A_1_block(2, 6);
    // A_1_block << 1, 0, -2, 0, 1, 0,
    //              0, 1, 0, -2, 0, 1;
    // for (int i = 0; i < 2 * n - 4; i += 2)
    //     A_1.block<2, 6>(i, i) = A_1_block;
    // // cout << A_1 << endl;

    // MatrixXd A_2_block(2, 4);
    // A_2_block << 1, 0, -1, 0,
    //              0, 1, 0, -1;
    // for (int i = 0; i < 2 * n - 2; i += 2)
    //     A_2.block<2, 4>(i, i) = A_2_block;
    /**********************************************************/
    SparseMatrix<double> A_1(2 * n - 4, 2 * n);
    SparseMatrix<double> A_2(2 * n - 2, 2 * n);
    MatrixXd A_3 = MatrixXd::Identity(2 * n, 2 * n);

    for (int i = 0; i < 2 * n - 4; i += 2) {
        A_1.insert(i, i) = 1;
        A_1.insert(i, i + 2) = -2;
        A_1.insert(i, i + 4) = 1;
        A_1.insert(i + 1, i + 1) = 1;
        A_1.insert(i + 1, i + 3) = -2;
        A_1.insert(i + 1, i + 5) = 1;
    }

    for (int i = 0; i < 2 * n - 2; i += 2) {
        A_2.insert(i, i) = 1;
        A_2.insert(i, i + 2) = -1;
        A_2.insert(i + 1, i + 1) = 1;
        A_2.insert(i + 1, i + 3) = -1;
    }

    hessian = 2 * (w_cost_smooth * A_1.transpose() * A_1 
                    + w_cost_length * A_2.transpose() * A_2 
                    + w_cost_ref * A_3);

    // 这里线性约束矩阵为单位矩阵
    for (int i = 0; i < 2 * n; i++) {
        linearMatrix.insert(i, i) = 1;
    }

    for (int i = 0; i < n; i++) {
        gradient[2*i] = this->ref_x[i];
        gradient[2*i+1] = this->ref_y[i];
        lowerBound[2*i] = gradient[2*i] + lb;
        lowerBound[2*i+1] = gradient[2*i+1] + lb;
        upperBound[2*i] = gradient[2*i] + ub;
        upperBound[2*i+1] = gradient[2*i+1] + ub;
    }

    gradient = -2 * w_cost_ref * gradient;

    // cout << "gradient: " << gradient.transpose();
    // cout << endl;
    // cout << "lowerBound: " << lowerBound.transpose();
    // cout << endl;
    // cout << "upperBound: " << upperBound.transpose();
    // cout << endl;
}


void QPSolver::initJerkMat(const vector<double> &l_min, const vector<double> &l_max, 
                            const double host_d1, const double host_d2, const double host_w, 
                            const SLPoint plan_start_sl,
                            const double w_cost_l,
                            const double w_cost_dl, 
                            const double w_cost_ddl, 
                            const double w_cost_dddl,
                            const double w_cost_center, 
                            const double w_cost_end_l,
                            const double w_cost_end_dl, 
                            const double w_cost_end_ddl) {
    int n = l_min.size();
    const double plan_start_l = plan_start_sl.l;
    const double plan_start_dl = plan_start_sl.l_prime;
    const double plan_start_ddl = plan_start_sl.l_dprime;
    // 这里的值分别为 线性约束矩阵的 行和列
    solver.data()->setNumberOfConstraints(9 * n -2);
    solver.data()->setNumberOfVariables(3 * n);

    this->hessian.resize(3 * n, 3 * n);
    this->gradient.resize(3 * n, 1);

    // 期望的终点状态
    double end_l_desire = 0;
    double end_dl_desire = 0;
    double end_ddl_desire = 0;

    /**
     * 计算 hessian
    */
    Eigen::SparseMatrix<double> H_L(3*n, 3*n);
    Eigen::SparseMatrix<double> H_DL(3*n, 3*n);
    Eigen::SparseMatrix<double> H_DDL(3*n, 3*n);
    Eigen::SparseMatrix<double> H_DDDL(n-1, 3*n);
    Eigen::SparseMatrix<double> H_L_END(3*n, 3*n);
    Eigen::SparseMatrix<double> H_DL_END(3*n, 3*n);
    Eigen::SparseMatrix<double> H_DDL_END(3*n, 3*n);
    // 生成H_L, H_DL, H_DDL, H_CENTER
    for (int i = 1; i <= n; i++) {
        H_L.insert(3 * i - 3, 3 * i - 3) = 1;
        H_DL.insert(3 * i - 2, 3 * i - 2) = 1;
        H_DDL.insert(3 * i - 1, 3 * i - 1) = 1;
    }
    Eigen::SparseMatrix<double> H_CENTER = H_L;
    // 生成 H_DDDL
    Eigen::MatrixXd H_DDDL_mat(n - 1, 3 * n);
    Eigen::MatrixXd H_dddl_sub(1, 6);
    H_dddl_sub << 0, 0, 1, 0, 0, -1;
    for (int i = 1; i < n; i++) {
        int row = i - 1;
        int col = 3 * i - 3;
        H_DDDL_mat.block<1, 6>(row, col) = H_dddl_sub;
    }
    H_DDDL = H_DDDL_mat.sparseView();

    // 生成 H_L_END H_DL_END H_DDL_END
    H_L_END.insert(3 * n - 3, 3 * n - 3) = 1;
    H_DL_END.insert(3 * n - 2, 3 * n - 2) = 1;
    H_DDL_END.insert(3 * n - 1, 3 * n - 1) = 1;

    // 生成二次规划的 H
    this->hessian = 2 * w_cost_l * (H_L.transpose() * H_L) +
                    w_cost_dl * (H_DL.transpose() * H_DL) + 
                    w_cost_ddl * (H_DDL.transpose() * H_DDL) +
                    // w_cost_dddl * (H_DDDL.transpose() * H_DDDL) +
                    w_cost_center * (H_CENTER.transpose() * H_CENTER) +
                    w_cost_end_l * (H_L_END.transpose() * H_L_END) +
                    w_cost_end_dl * (H_DL_END.transpose() * H_DL_END) +
                    w_cost_end_ddl * (H_DDL_END.transpose() * H_DDL_END);
    /**
     * 生成 F，一次向量矩阵
     * 这里只对 l 有约束
    */
    vector<double> center_line(3 * n, 0);
    for (int i = 0; i < l_min.size(); i++) {
        center_line[i] = (l_min[i] + l_max[i]) / 2.0;
    }
    Eigen::SparseMatrix<double> f(3 * n, 1);
    for (int i = 1; i <= n; i++) {
        f.insert(3 * i - 3, 0) = -2 * center_line[i-1]; 
    }
    f = w_cost_center * f;
    // f(3 * n - 3) = f(3 * n - 3) - 2 * end_l_desire * w_cost_end_l;
    // f(3 * n - 2) = f(3 * n - 2) - 2 * end_dl_desire * w_cost_end_dl;
    // f(3 * n - 1) = f(3 * n - 1) - 2 * end_ddl_desire * w_cost_end_ddl;
    this->gradient = f;

    // 9n-2 = (2n-2) + (4n) + 3n
    Eigen::MatrixXd linear(9*n-2, 3*n);
    Eigen::VectorXd lowerB(9*n - 2);
    Eigen::VectorXd upperB(9*n - 2);

    /**
     * 等式约束
    */
    double ds = 0.5;
    Eigen::MatrixXd Aeq_mat(2*n-2, 3*n);
    Eigen::MatrixXd Aeq_sub(2, 6);
    Eigen::VectorXd beq = Eigen::VectorXd::Zero(2*n-2);
    Aeq_sub << 1, ds, pow(ds, 2) / 3, -1, 0, pow(ds, 2) / 6,
               0, 1,  ds / 2,          0, -1, ds / 2;

    for (int i = 1; i < n; i++) {
        int row = 2 * i - 2;
        int col = 3 * i - 3;
        Aeq_mat.block<2, 6>(row, col) = Aeq_sub;
    }


    /**
     * 不等式约束
    */
    // Eigen::SparseMatrix<double> A(4*n, 3*n);
    // A
    double d1 = host_d1;
    double d2 = host_d2;
    double w = host_w;
    Eigen::MatrixXd A_mat(4*n, 3*n);
    Eigen::MatrixXd A_sub(4, 3);
    A_sub << 1, d1, 0,
             1, d1, 0,
             1,-d2, 0,
             1,-d2, 0;
    for (int i = 1; i <= n; i++) {
        int row = 4 * i - 4;
        int col = 3 * i - 3;
        A_mat.block<4, 3>(row, col) = A_sub;
    }


    // B
    Eigen::VectorXd b_mat_lb(4 * n);
    Eigen::VectorXd b_mat_ub(4 * n);
    int front_index = ceil(d1 / ds);
    int back_index = ceil(d2 / ds);
    for (int i = 0; i < n; i++) {
        // 左前 右前 的index
        int index1 = min(i + front_index, n - 1);
        // 左后 右后的index
        int index2 = max(i - back_index, 0);
        Eigen::VectorXd b_sub_ub(4);
        Eigen::VectorXd b_sub_lb(4);
        b_sub_ub << l_max[index1] - w / 2,
                    l_max[index1] + w / 2,
                    l_max[index2] - w / 2,
                    l_max[index2] + w / 2;
        b_sub_lb << l_min[index1] - w / 2,
                    l_min[index1] + w / 2,
                    l_min[index2] - w / 2,
                    l_min[index2] + w / 2;
        b_mat_lb.segment(4 * i, 4) = b_sub_lb;
        b_mat_ub.segment(4 * i, 4) = b_sub_ub;
    }

    // 生成 lb ub 主要是对规划起点做约束
    Eigen::MatrixXd identity = Eigen::MatrixXd::Identity(3*n, 3*n);
    Eigen::VectorXd lb(3*n);
    Eigen::VectorXd ub(3*n);
    lb.fill(-OsqpEigen::INFTY);
    ub.fill(OsqpEigen::INFTY);
    lb(0) = ub(0) = plan_start_l;
    lb(1) = ub(1) = plan_start_dl;
    lb(2) = ub(2) = plan_start_ddl;

    // 将等式约束、不等式约束、范围约束组合起来
    linear << Aeq_mat,
              A_mat,
              identity;
    lowerB << beq,
              b_mat_lb,
              lb;
    upperB << beq,
              b_mat_ub,
              ub;
    this->linearMatrix = linear.sparseView();
    this->lowerBound = lowerB;
    this->upperBound = upperB;
}


void QPSolver::initFEMSmoothMat() {
    int NumberOfVariables = 2;   //A矩阵的列数
    int NumberOfConstraints = 3; //A矩阵的行数
    solver.data()->setNumberOfVariables(NumberOfVariables);
    solver.data()->setNumberOfConstraints(NumberOfConstraints);

    hessian.resize(NumberOfVariables, NumberOfVariables);
    gradient.resize(NumberOfVariables);
    linearMatrix.resize(NumberOfConstraints, NumberOfVariables);
    lowerBound.resize(NumberOfConstraints);
    upperBound.resize(NumberOfConstraints);

    hessian.insert(0, 0) = 1;
    hessian.insert(1, 0) = -1;
    hessian.insert(0, 1) = -1;
    hessian.insert(1, 1) = 2;
    std::cout << "hessian:" << std::endl
                << hessian << std::endl;

    gradient << -2, -6;
    std::cout << "gradient:" << std::endl
                << gradient << std::endl;
    
    linearMatrix.insert(0, 0) = 1;
    linearMatrix.insert(0, 1) = 1;
    linearMatrix.insert(1, 0) = -1;
    linearMatrix.insert(1, 1) = 2;
    linearMatrix.insert(2, 0) = 2;
    linearMatrix.insert(2, 1) = 1;
    std::cout << "linearMatrix:" << std::endl
                << linearMatrix << std::endl;
    // 约束下界，没有就用负无穷表示
    lowerBound << -OsqpEigen::INFTY, -OsqpEigen::INFTY, -OsqpEigen::INFTY;
    // std::cout << "lowerBound:" << std::endl
    //             << lowerBound << std::endl;

    upperBound << 2, 2, 3;
    // std::cout << "upperBound:" << std::endl
    //             << upperBound << std::endl;
}

bool QPSolver::checkMat() {
    if (!solver.data()->setHessianMatrix(hessian)) return false;
    if (!solver.data()->setGradient(gradient)) return false;
    if (!solver.data()->setLinearConstraintsMatrix(linearMatrix)) return false;
    if (!solver.data()->setLowerBound(lowerBound)) return false;
    if (!solver.data()->setUpperBound(upperBound)) return false;
    return true;
}

bool QPSolver::run(vector<Point> &smooth_points) {
    solver.clearSolver();
    solver.data()->clearHessianMatrix();
    solver.data()->clearLinearConstraintsMatrix();

    vector<double> smooth_x, smooth_y;          // 平滑后的轨迹坐标
    // vector<Point> smooth_points;
    
    if (checkMat()) {
        if (!solver.initSolver()) {
            cout << "Can't init solver !!!" << endl;
            return false;
        }
        
    } else {
        cout << "initilize QP solver failed" << endl;
        return false;
    }

    Eigen::VectorXd QPSolution;
    clock_t time_start, time_end;
    time_start = clock();
    this->solver.solveProblem();
    QPSolution = this->solver.getSolution();
    time_end = clock();

    
    for (int i = 0; i < QPSolution.rows(); i += 2) {
        Point p(QPSolution[i], QPSolution[i+1]);
        smooth_points.push_back(p);
    }
    
    if (this->show) {
        std::cout << "time use:" << 1000 * (time_end - time_start) / (double)CLOCKS_PER_SEC << "ms" << std::endl;
        cout << "QPSolution: " << endl << QPSolution.transpose() << endl;
        
        for (int i = 0; i < QPSolution.rows(); i += 2) {
            // cout << QPSolution.row(i);
            smooth_x.push_back(QPSolution[i]);
            smooth_y.push_back(QPSolution[i+1]);
        }
        plt::plot(ref_x,ref_y,"b--");
        plt::plot(smooth_x, smooth_y,"r" "o");
        plt::plot(smooth_x, smooth_y,"y");
        plt::grid(true);
        plt::xlim(0, 10);
        plt::show();
    }

    return true;
}


bool QPSolver::runJark(vector<double> &path_l, vector<double> &path_dl, vector<double> &path_ddl) {
    solver.clearSolver();
    solver.data()->clearHessianMatrix();
    solver.data()->clearLinearConstraintsMatrix();

    vector<double> smooth_x, smooth_y;          // 平滑后的轨迹坐标
    // vector<Point> smooth_points;
    
    if (checkMat()) {
        if (!solver.initSolver()) {
            cout << "Can't init solver !!!" << endl;
            return false;
        } 
    } else {
        cout << "initilize QP solver failed" << endl;
        return false;
    }

    Eigen::VectorXd QPSolution;
    clock_t time_start, time_end;
    time_start = clock();
    this->solver.solveProblem();
    QPSolution = this->solver.getSolution();
    time_end = clock();

    
    for (int i = 0; i < QPSolution.rows(); i += 3) {
        path_l.push_back(QPSolution[i]);
        path_dl.push_back(QPSolution[i+1]);
        path_ddl.push_back(QPSolution[i+2]);
    }
    return true;
}


bool QPSolver::runQuinticPoly(vector<double> &params) {
    solver.clearSolver();
    solver.data()->clearHessianMatrix();
    solver.data()->clearLinearConstraintsMatrix();
    
    if (checkMat()) {
        if (!solver.initSolver()) {
            cout << "Can't init solver !!!" << endl;
            return false;
        }
        
    } else {
        cout << "initilize QP solver failed" << endl;
        return false;
    }

    Eigen::VectorXd QPSolution;
    clock_t time_start, time_end;
    time_start = clock();
    this->solver.solveProblem();
    QPSolution = this->solver.getSolution();
    time_end = clock();
    std::cout << "time use:" << 1000 * (time_end - time_start) / (double)CLOCKS_PER_SEC << "ms" << std::endl;
    for (int i = 0; i < QPSolution.rows(); i ++) {
        params.push_back(QPSolution[i]);
    }
    return true;
}
// int main() {
//     vector<double> ref_x_1 = {0, 1, 2, 3, 4, 5, 5.5, 6, 6.5, 7};
//     vector<double> ref_y_1 = {0, 0, 2, 4, 4, 4.5, 5, 6, 9, 12};
//     vector<double> ref_x_2 = {0, 1, 2, 3, 4, 5, 5.5, 6, 6.5, 7};
//     vector<double> ref_y_2 = {4, 3, 2.5, 4, 4, 4.5, 5, 6, 9, 12};

//     QPSolver qp_solver(false);
//     qp_solver.updateRefePoint(ref_x_1, ref_y_1);
//     qp_solver.initFEMSmoothMat(10, 2, 3, 2, -0.3, 0.3);
//     qp_solver.run();

//     qp_solver.updateRefePoint(ref_x_2, ref_y_2);
//     qp_solver.initFEMSmoothMat(10, 2, 3, 2, -0.3, 0.3);
//     qp_solver.run();
// }